#ifdef ENABLE_MAPLE_SAN

#include "asan_config.h"
#include "asan_asap.h"
#include "san_common.h"
#include "asan_razor.h"

namespace maple {

double ASAPThreshold;
std::vector<ASAPStmtKeyCountPair> OrderedStmtKeyCountPairs;
std::set<ASAPStmtKey> ASAPAsanRelevantStmtKey;
std::map<const ASAPStmtKey, AsanBlockCost> ASAPStmtKey2CostMap;
std::set<ASAPStmtKey> ASAPToRemoveAsanIDs;
std::set<StmtID> ASAPToRemoveStmtIDs;

void InitializeOrderedStmtKeyCountPairs() {
  OrderedStmtKeyCountPairs.clear();
  std::map<ASAPStmtKey, AsanRazorCheckCount> tmp;
  for (auto &pair : SanRazorStmtCountMap) {
    ASAPStmtKey key(pair.first.GetStmtID(), pair.first.src_path);
    if (tmp.find(key) != tmp.end()) {
      tmp[key] += pair.second;
    } else {
      tmp[key] = pair.second;
    }
    if (SanRazorAsanRelevantStmtKey.find(pair.first) != SanRazorAsanRelevantStmtKey.end()) {
      ASAPAsanRelevantStmtKey.insert(key);
      ASAPStmtKey2CostMap[key] = 1.0;
    } else {
      ASAPStmtKey2CostMap[key] += 1.0;
    }
  }
  for (auto &pair : tmp) {
    OrderedStmtKeyCountPairs.push_back(pair);
  }
  // sort in reverse order
  std::sort(OrderedStmtKeyCountPairs.begin(), OrderedStmtKeyCountPairs.end(),
    [](const ASAPStmtKeyCountPair &a, const ASAPStmtKeyCountPair &b) { return a.second > b.second; });
}

void LoadASAPThresholdNumMode() {
  char *threshold = std::getenv("ASAP_THRESHOLD");
  if (threshold == nullptr) {
    ASAPThreshold = ASAN_ASAP_DEFAULT_ASAN_CHECK_NUM_THRESHOLD;
  } else {
    ASAPThreshold = std::atof(threshold);
  }
}

void LoadASAPThresholdCostMode() {
  char *threshold = std::getenv("ASAP_THRESHOLD");
  if (threshold == nullptr) {
    ASAPThreshold = ASAN_ASAP_DEFAULT_ASAN_CHECK_COST_THRESHOLD;
  } else {
    ASAPThreshold = std::atof(threshold);
  }
}

void LoadASAPThreshold(ASAPMode mode) {
  switch (mode) {
    case ASAPMode::NUM:
      LoadASAPThresholdNumMode();
      break;
    case ASAPMode::COST:
      LoadASAPThresholdCostMode();
      break;
    default:
      LogInfo::MapleLogger() << "Invalid ASAP mode\n";
      break;
  }
}

void DecideToBeRemovedNumMode() {
  // Note this function removes the asan checks of the whole program,
  // rather than a single source file
  size_t num_asan_checks = 0;
  for (auto &pair : OrderedStmtKeyCountPairs) {
    if (ASAPAsanRelevantStmtKey.find(pair.first) == ASAPAsanRelevantStmtKey.end()) {
      continue;
    }
    num_asan_checks++;
  }
  size_t ASAPThreshold_abs = (size_t)(num_asan_checks * ASAPThreshold);
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << "To remove following checks:\n" << std::hex;
  }
  size_t del_sum = 0;
  size_t idx = 0;
  while (del_sum <= ASAPThreshold_abs && idx < OrderedStmtKeyCountPairs.size()) {
    auto& pair = OrderedStmtKeyCountPairs[idx];
    // Here we merely consider the asan relevant stmts
    if (ASAPAsanRelevantStmtKey.find(pair.first) == ASAPAsanRelevantStmtKey.end()) {
      idx++;
      continue;
    }
    del_sum++;
    idx++;
    ASAPToRemoveAsanIDs.insert(pair.first);
    if (ASAN_DEBUG_MODE) {
      LogInfo::MapleLogger() << pair.first.src_path << " " << pair.first.stmtid << " " << pair.second << "\n";
    }
  }
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << std::dec << "To delete " << ASAPToRemoveAsanIDs.size() << "(" << \
      ((double)ASAPToRemoveAsanIDs.size() / (double)num_asan_checks) << ")" << \
      " asan checks\n";
  }
}

void DecideToBeRemovedCostMode() {
  // Note this function removes the asan checks of the whole program,
  // rather than a single source file
  long double total = 0.0;
  size_t num_asan_checks = 0;
  for (auto &pair : OrderedStmtKeyCountPairs) {
    if (ASAPAsanRelevantStmtKey.find(pair.first) == ASAPAsanRelevantStmtKey.end()) {
      continue;
    }
    num_asan_checks++;
    if (ASAPStmtKey2CostMap.find(pair.first) == ASAPStmtKey2CostMap.end()) {
      total += pair.second;
    } else {
      total += pair.second * ASAPStmtKey2CostMap[pair.first];
    }
  }
  size_t idx = 0;
  long double del_sum = 0.0;
  double ASAPThreshold_abs = total * ASAPThreshold;
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << "To remove following checks:\n" << std::hex;
  }
  while (del_sum < ASAPThreshold_abs && idx < OrderedStmtKeyCountPairs.size()) {
    auto& pair = OrderedStmtKeyCountPairs[idx];
    // Here we merely consider the asan relevant stmts
    if (ASAPAsanRelevantStmtKey.find(pair.first) == ASAPAsanRelevantStmtKey.end()) {
      idx++;
      continue;
    }
    if (ASAPStmtKey2CostMap.find(pair.first) == ASAPStmtKey2CostMap.end()) {
      del_sum += pair.second;
    } else {
      del_sum += pair.second * ASAPStmtKey2CostMap[pair.first];
    }
    idx++;
    ASAPToRemoveAsanIDs.insert(pair.first);
    if (ASAN_DEBUG_MODE) {
      LogInfo::MapleLogger() << pair.first.src_path << " " << pair.first.stmtid << " " << pair.second << "\n";
    }
  }
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << std::dec << "To delete " << ASAPToRemoveAsanIDs.size() << "(" << \
      ((double)ASAPToRemoveAsanIDs.size() / (double)num_asan_checks) << ")" << \
      " asan checks\n";
  }
}

void DecideToBeRemoved(ASAPMode mode) {
  switch (mode) {
    case ASAPMode::NUM:
      DecideToBeRemovedNumMode();
      break;
    case ASAPMode::COST:
      DecideToBeRemovedCostMode();
      break;
    default:
      LogInfo::MapleLogger() << "Invalid ASAP mode\n";
      break;
  }
}

void SetThisFileASAPToRemoveStmtIDs(std::string filename) {
  for (auto &key : ASAPToRemoveAsanIDs) {
    if (key.src_path == filename) {
      ASAPToRemoveStmtIDs.insert(key.stmtid);
    }
  }
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << "ASAP will remove " << std::dec << \
      ASAPToRemoveStmtIDs.size() << " asan checks of " << filename << "\n";
  }
}

void ASAPRemoveAsanChecks(std::string filename, BlockNode* block) {
  std::vector<StmtBlockNodePair> stmtBlockList = getAllOrderedStmtNodeList(block);
  if (ASAN_DEBUG_MODE) {
    // block->Dump(0);
    LogInfo::MapleLogger() << "Before ASAP " << std::dec << \
      stmtBlockList.size() << " StmtNodes\n";
  }
  std::set<StmtNode*> toRemoveStmtNodeSet;
  size_t num_asan_blocks = 0;
  for (auto &asanBlock : *AsanBlockListPtr->blocks) {
    StmtID first_id = asanBlock.first->GetStmtID();
    if (ASAPToRemoveStmtIDs.find(first_id) != ASAPToRemoveStmtIDs.end()) {
      num_asan_blocks++;
      StmtNode* cur = asanBlock.first;
      while (cur != nullptr && cur != asanBlock.last) {
        toRemoveStmtNodeSet.insert(cur);
        cur = cur->GetNext();
      }
      toRemoveStmtNodeSet.insert(asanBlock.last);
    }
  }
  if (ASAN_DEBUG_MODE) {
    LogInfo::MapleLogger() << "ASAP removes " << std::dec << \
      num_asan_blocks << " asan blocks " << \
      toRemoveStmtNodeSet.size() << " StmtNodes\n";
  }
  for (auto &pair : stmtBlockList) {
    if (toRemoveStmtNodeSet.find(pair.first) != toRemoveStmtNodeSet.end()) {
      pair.second->RemoveStmt(pair.first);
    }
  }
  if (ASAN_DEBUG_MODE) {
    stmtBlockList = getAllOrderedStmtNodeList(block);
    LogInfo::MapleLogger() << "After ASAP " << std::dec << \
      stmtBlockList.size() << " StmtNodes\n";
    // block->Dump(0);
  }
}

} // namespace maple

#endif // ENABLE_MAPLE_SAN